# -*- coding: utf-8 -*-
"""
Created on Fri Nov 27 10:33:29 2020

@author: zhaodi
"""

import numpy as np


def czt(x, m=None, w=None, a=None):
    n = len(x)

    if m is None: m = n
    if w is None: w = np.exp(-2j * np.pi / m)
    if a is None: a = 1

    wExponents1 = np.arange(0, n) ** 2 / 2.0  # n elements         [ 0 .. (n - 1) ]
    wExponents2 = - np.arange(-(n - 1), m) ** 2 / 2.0  # m + n - 1 elements [ -(n - 1) .. +(m - 1) ]
    wExponents3 = np.arange(0, m) ** 2 / 2.0  # m elements         [ 0        ..  (m - 1) ]

    xx = x * a ** -np.arange(n) * w ** wExponents1

    # Determine next-biggest FFT of power-of-two

    nfft = 1
    while nfft < (m + n - 1):
        nfft += nfft

    fxx = np.fft.fft(xx, nfft)

    ww = w ** wExponents2

    fww = np.fft.fft(ww, nfft)

    fyy = fxx * fww

    yy = np.fft.ifft(fyy, nfft)

    # select output

    yy = yy[n - 1: m + n - 1]

    y = yy * w ** wExponents3

    return y


def invczt(X, F, T1, T2, n, flag, s):
    #    fn = (np.linspace(0,n-1,n-1))/n
    #    fs = 1/(T2-T1)
    #    fy = fn*fs
    delta_F = F[1] - F[0]
    delta_t = (T2 - T1) / n
    w = np.exp(-1 * 1j * 2 * np.pi * delta_t * delta_F)
    a = np.exp(1j * 2 * np.pi * T1 * delta_F)
    t = np.linspace(T1, T2, n)
    win = 0.5 - s * np.cos(2 * np.pi / len(X) * (np.linspace(0, len(X) - 1, len(X))))
    if flag == 1:
        X = X * win
    Y = 2 * np.conj(czt(np.conj(X), n, w, a)) * np.exp(1j * 2 * np.pi * F[0] * t) / len(X)

    return Y, t


def pmusic(x, p):
    m = 10
    s = np.matrix(np.ones([m - 1, len(x) - m]))
    s = np.row_stack((x[0:len(x) - m], s))
    for i in range(1, m):
        s[i, :] = x[i:len(x) - m + i]
    [u, d, v] = np.linalg.svd(s)
    un = u[:, p:m]
    res = np.fft.ifft(un.T, 1024)
    res = 1 / np.array(sum(abs(res)))
    return (res)


def remove_back(dataMetal, dataAir):
    # dataMetal = np.mean(data.real.T, 1)
    dataRemoveBac = np.mean(dataMetal.T - np.mean(dataAir.T, 1)[:, np.newaxis], 1)
    dataRemoveBac = np.abs(dataRemoveBac)
    return list(dataRemoveBac).index(max(dataRemoveBac))

